#include <iostream>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdarg.h>
#include <time.h>
#include <unistd.h>
#include <string.h>
#include <sys/resource.h>
#include <sys/wait.h>
#include <sys/stat.h>
#include <sys/user.h>
#include <sys/ptrace.h>
#include <sys/types.h>
#include <sys/fcntl.h>
#include <seccomp.h>
#include <sys/syscall.h>
#include "table.h"
typedef unsigned int uint;
typedef long long LL;
typedef unsigned char uchar;
const int BUF_SIZE = 1024*1024*5;
char buf[BUF_SIZE+3];

//define judge status
const int OJ_OK = 0;
const int OJ_ML = 1;
const int OJ_TL = 2;
const int OJ_RE = 3;
const int OJ_OL = 4;
const int OJ_DC = 5; //dangerous calls
const int OJ_CE = 6;
const int OJ_WA = 7;
const int OJ_PE = 8;

const int COMPILE = 0;
const int RUN = 1;
const int RUN_COMMAND = 3;

//定义语言
const int C = 1;
const int CPP = 2;
const int JAVA = 3;
const int PASCAL = 4;
const int PYTH = 5;
const int CPP_SPJ = 6;
//判题配置
const bool debug = false;
const char* logPath = "/home/judger/logs/JudgeClient.log"; //"/var/log/JudgeClient.log";
const char* ojUrl = "localhost:8080/OnlineJudge"; //oj项目地址
const char* clientName = "local1"; //判题机账号
const char* clientPassword = "haha"; //判题机密码
const int customType = 4;
const int customTimeLimit = 5; //用户自测试时间设为5s
const int customMemoryLimit = 256; //用户代码占用内存设为256M

std::string getCodeSuffix(int lang) {
    switch(lang) {
        case C:
            return "c";
        case CPP:
            return "cpp";
        case JAVA:
            return "java";
        case PASCAL:
            return "pascal";
        case PYTH:
            return "py";
    }
    return "";
}

void putLog(const char* fmt, ...) {
    time_t t;
    time(&t);
    tm ttime = *localtime(&t);
    sprintf(buf, "%4d-%02d-%02d %02d:%02d:%02d log: ", ttime.tm_year + 1900, ttime.tm_mon + 1, ttime.tm_mday, ttime.tm_hour, ttime.tm_min, ttime.tm_sec);
    std::string llog = (std::string)buf;
    va_list args;
    va_start(args, fmt);
    vsprintf(buf, fmt, args);
    va_end(args);
    llog += (std::string)buf;
    FILE* logFile = fopen(logPath, "a+");
    if (logFile == NULL) {
        printf("open logfile error.\n");
        return;
    }
    fputs(llog.c_str(), logFile);
    fclose(logFile);
}

LL getFileSize(char* file) {
    struct stat f_stat;
    if (stat(file, &f_stat) == -1) {
        return 0;
    }
    return (LL) f_stat.st_size;
}

int exeCmd(const char* fmt, ...) {
    va_list args;
    va_start(args, fmt);
    vsprintf(buf, fmt, args);
    va_end(args);
    int ret = system(buf);
    return ret;
}

FILE* getCmdOutput(const char* fmt, ...) {
    FILE* ret = NULL;
    va_list args;
    va_start(args, fmt);
    vsprintf(buf, fmt, args);
    va_end(args);
    ret = popen(buf, "r");
    return ret;
}

bool isBlank(char* str) {
    char* buf = str;
    if (str == NULL) return true;
    while (*buf != '\0') {
        if (*buf != ' ') {
            return false;
        }
    }
    return true;
}

void EXIT_PROC(int code, const char* msg) {
#ifdef LOG
    putLog("%s with exit code %d\n", msg, code);
#endif
    if (debug) {
        perror(msg);
    }
    exit(code);
}

void initCtx(int act, int lang) {
    int ret;
    scmp_filter_ctx ctx = NULL;
    if(act == COMPILE) {
        int *table = BLACK_LIST_COMPILE;
        ctx = seccomp_init(SCMP_ACT_ALLOW);
        if (ctx == NULL) {
            EXIT_PROC(-1, "seccomp_init error");
        }
        for (int i=0; table[i] != -1; ++i) {
            ret = seccomp_rule_add(ctx, SCMP_ACT_KILL, table[i], 0);
            if (ret != 0) {
                EXIT_PROC(-1, "compile black list rule add error");
            }
        }
    }
    else {
        if (lang == C || lang == CPP || lang == CPP_SPJ) {
            int *table = WHITE_LIST_CPP;
            ctx = seccomp_init(SCMP_ACT_KILL);
            if (ctx == NULL) {
                EXIT_PROC(-1, "seccomp_init error");
            }
            for (int i=0; table[i] != -1; ++i) {
                ret = seccomp_rule_add(ctx, SCMP_ACT_ALLOW, table[i], 0);
                if (ret != 0) {
                    EXIT_PROC(-1, "white list rule add error");
                }
            }
            ret = seccomp_rule_add(ctx, SCMP_ACT_ALLOW, SYS_execve, 1, SCMP_CMP(0, SCMP_CMP_EQ, (scmp_datum_t)"./Main"));
            if (ret != 0) {
                EXIT_PROC(-1, "execve rule add error");
            }
            // forbidden the write for open and openat
            ret = seccomp_rule_add(ctx, SCMP_ACT_ALLOW, SYS_openat, 1, SCMP_CMP(2, SCMP_CMP_MASKED_EQ, O_WRONLY|O_RDWR, 0));
            if (ret != 0) {
                EXIT_PROC(-1, "openat rule add error");
            }
            ret = seccomp_rule_add(ctx, SCMP_ACT_ALLOW, SYS_open, 1, SCMP_CMP(1, SCMP_CMP_MASKED_EQ, O_WRONLY|O_RDWR, 0));
            if (ret != 0) {
                EXIT_PROC(-1, "open rule add error");
            }
            if (lang == CPP_SPJ) {
                ret = seccomp_rule_add(ctx, SCMP_ACT_ALLOW, SYS_execve, 1, SCMP_CMP(0, SCMP_CMP_EQ, (scmp_datum_t)"./spj.out"));
                if (ret != 0) {
                    EXIT_PROC(-1, "execve rule add error");
                }
            }
        }
        else {
            int *table = BLACK_LIST_GENERAL;
            ctx = seccomp_init(SCMP_ACT_ALLOW);
            if (ctx == NULL) {
                EXIT_PROC(-1, "seccomp_init error");
            }
            for (int i=0; table[i] != -1; ++i) {
                ret = seccomp_rule_add(ctx, SCMP_ACT_KILL, table[i], 0);
                if (ret != 0) {
                    EXIT_PROC(-1, "black list rule add error");
                }
            }
            if (lang == JAVA) {
                ret = seccomp_rule_add(ctx, SCMP_ACT_KILL, SYS_execve, 1, SCMP_CMP(0, SCMP_CMP_NE, (scmp_datum_t)"/usr/bin/java"));
                if (ret != 0) {
                    EXIT_PROC(-1, "execve rule add error");
                }
            }
            //copy from QingdaoU judger for socket
            ret = seccomp_rule_add(ctx, SCMP_ACT_ERRNO(EACCES), SCMP_SYS(socket), 0);
            if (ret != 0) {
                EXIT_PROC(-1, "socket rule add error");
            }
            //java use policy to control file
            if (lang != JAVA) {
                // forbidden the write for openat
                ret = seccomp_rule_add(ctx, SCMP_ACT_KILL, SYS_openat, 1, SCMP_CMP(2, SCMP_CMP_MASKED_EQ, O_WRONLY, O_WRONLY));
                if (ret != 0) {
                    EXIT_PROC(-1, "openat rule add error");
                }
                ret = seccomp_rule_add(ctx, SCMP_ACT_KILL, SYS_openat, 1, SCMP_CMP(2, SCMP_CMP_MASKED_EQ, O_RDWR, O_RDWR));
                if (ret != 0) {
                    EXIT_PROC(-1, "openat rule add error");
                }
                // forbidden the write for open
                ret = seccomp_rule_add(ctx, SCMP_ACT_KILL, SYS_open, 1, SCMP_CMP(1, SCMP_CMP_MASKED_EQ, O_WRONLY, O_WRONLY));
                if (ret != 0) {
                    EXIT_PROC(-1, "open rule add error");
                }
                ret = seccomp_rule_add(ctx, SCMP_ACT_KILL, SYS_open, 1, SCMP_CMP(1, SCMP_CMP_MASKED_EQ, O_RDWR, O_RDWR));
                if (ret != 0) {
                    EXIT_PROC(-1, "open rule add error");
                }
            }
            // forbidden change default signal hander SIGCHLD, SIGXCPU, SIGQUIT, SIGSYS, SIGXFSZ
            ret = seccomp_rule_add(ctx, SCMP_ACT_KILL, SYS_rt_sigaction, 1, SCMP_CMP(0, SCMP_CMP_EQ, SIGCHLD));
            if (ret != 0) {
                EXIT_PROC(-1, "sigaction rule add error");
            }
            ret = seccomp_rule_add(ctx, SCMP_ACT_KILL, SYS_rt_sigaction, 1, SCMP_CMP(0, SCMP_CMP_EQ, SIGXCPU));
            if (ret != 0) {
                EXIT_PROC(-1, "sigaction rule add error");
            }
            if (lang != JAVA) { //signal SIGQUIT and SIGXFSZ is useful for JVM
                ret = seccomp_rule_add(ctx, SCMP_ACT_KILL, SYS_rt_sigaction, 1, SCMP_CMP(0, SCMP_CMP_EQ, SIGQUIT));
                if (ret != 0) {
                    EXIT_PROC(-1, "sigaction rule add error");
                }
                ret = seccomp_rule_add(ctx, SCMP_ACT_KILL, SYS_rt_sigaction, 1, SCMP_CMP(0, SCMP_CMP_EQ, SIGXFSZ));
                if (ret != 0) {
                    EXIT_PROC(-1, "sigaction rule add error");
                }
            }
            ret = seccomp_rule_add(ctx, SCMP_ACT_KILL, SYS_rt_sigaction, 1, SCMP_CMP(0, SCMP_CMP_EQ, SIGSYS));
            if (ret != 0) {
                EXIT_PROC(-1, "sigaction rule add error");
            }
        }
    }
    //将过滤规则加入内核
    if (seccomp_load(ctx) != 0) {
        EXIT_PROC(-1, "seccomp load error");
    }
    //过滤规则加入内核之后释放ctx的空间，加入内核之后release不改变规则
    seccomp_release(ctx);
}

//time_lmt以s计，mem_lmt以mb计
//运行时间以ms计, 运行内存以mb计
//子程序的运行函数
int runChild(int act, char* path, char** command, char* workdir, int lang,
        const int& time_lmt, const int& mem_lmt) {

    int pid = getpid();
    while (setpgid(pid, pid) != 0) usleep(50000);

    //切换到判题目录
    chdir(workdir);
    freopen("data.in", "r", stdin);
    freopen("user.out", "w", stdout);
    freopen("user.err", "w", stderr);
    //设置到判题用户judger中
    while (setgid(6666) != 0) usleep(50000);
    while (setuid(6666) != 0) usleep(50000);
    while (setresuid(6666, 6666, 6666) != 0) usleep(50000);

    rlimit limit;
    limit.rlim_cur = time_lmt + 1;//限制CPU运行时间
    limit.rlim_max = limit.rlim_cur;
    setrlimit(RLIMIT_CPU, &limit);
    //限制文件输出大小
    limit.rlim_max = limit.rlim_cur = (64<<20);
    setrlimit(RLIMIT_FSIZE, &limit);
    if (lang == JAVA) {
        //限制进程数
        limit.rlim_max = limit.rlim_cur = 22;
        setrlimit(RLIMIT_NPROC, &limit);
        //java使用-X选项限制内存大小
    }
    else if (act == COMPILE) {
        //编译需要的进程数多一些
        limit.rlim_max = limit.rlim_cur = 8;
        setrlimit(RLIMIT_NPROC, &limit);
        //限制堆栈和内存
        limit.rlim_max = limit.rlim_cur = ((mem_lmt*2)<<20);
        setrlimit(RLIMIT_AS, &limit);
        setrlimit(RLIMIT_DATA, &limit);
        setrlimit(RLIMIT_STACK, &limit);
    }
    else {
        //限制进程数
        limit.rlim_max = limit.rlim_cur = 2;
        setrlimit(RLIMIT_NPROC, &limit);
        //限制堆栈和内存
        limit.rlim_max = limit.rlim_cur = ((mem_lmt*2)<<20);
        setrlimit(RLIMIT_AS, &limit);
        setrlimit(RLIMIT_DATA, &limit);
        setrlimit(RLIMIT_STACK, &limit);
    }
    //加载过滤规则
    initCtx(act, lang);
    execv(path, command);
    exit(0);
}

int watchChild(int pid, const int& time_lmt,
        const int& mem_lmt, double& time, double& mem, char* workdir) {
    int tmem, status, ttime;
    int ret = OJ_OK;
    rusage rusa;
    int tt = time_lmt*2 + 1;
    while(1) {
        if (wait4(pid, &status, WNOHANG, &rusa) > 0) {
            break;
        }
        sleep(1);
        --tt;
        if (tt == 0) {
            kill(-pid, SIGKILL);
            usleep(2000);
        }
    }
    //wait4(pid, &status, 0, &rusa);

    tmem = rusa.ru_maxrss;
    ttime = rusa.ru_utime.tv_sec * 1000 + rusa.ru_utime.tv_usec / 1000;
    ttime += rusa.ru_stime.tv_sec * 1000 + rusa.ru_stime.tv_usec / 1000;

    //如果运行程序为正常终止
    if (WIFEXITED(status)) {
        int t = WEXITSTATUS(status);
        if (t != 0) ret = OJ_RE;
    }
    //如果运行程序被信号终止
    else if (WIFSIGNALED(status)) {
        int t = WTERMSIG(status);
        switch (t) {
            case SIGCHLD:
            case SIGXCPU:
            case SIGQUIT:
                ret = OJ_TL;
                break;
            case SIGXFSZ:
                ret = OJ_OL;
                break;
            case SIGSYS:
                ret = OJ_DC;
                break;
            default:
                ret = OJ_RE;
        }
    }
    if (ttime > time_lmt*1000) {
        ret = OJ_TL;
    }
    else if (tmem > (mem_lmt<<10)) {
        ret = OJ_ML;
    }
    time = ttime;
    mem = (double)tmem/1024.0;
    return ret;
}

int runCommand(int act, char* path, char** command, char* workdir, int lang,
        const int& time_lmt, const int& mem_lmt, double& time, double& mem)
{
    int ret = -1;
    int pid = fork();
    if (pid == 0) {
        runChild(act, path, command, workdir, lang, time_lmt, mem_lmt);
    }
    else {
        ret = watchChild(pid, time_lmt, mem_lmt, time, mem, workdir);
    }
    return ret;
}

int runSPJ(int lang, int time_lmt, char* workdir) {
    int pid = fork();
    if (pid == 0) {
        //切换到判题目录
        chdir(workdir);
        int pid = getpid();
        while (setpgid(pid, pid) != 0) usleep(50000);
        //设置到判题用户judger中
        while (setgid(6666) != 0) usleep(50000);
        while (setuid(6666) != 0) usleep(50000);
        while (setresuid(6666, 6666, 6666) != 0) usleep(50000);
        //加载过滤规则
        initCtx(RUN, CPP_SPJ);
        if (lang == C) {
            const char* command[] = {"./spj.out", "./data.in", "./data.out", "./Main.c", "user.out", NULL};
            execv("./spj.out", (char**)command);
        }
        else if (lang == CPP) {
            const char* command[] = {"./spj.out", "./data.in", "./data.out", "./Main.cpp", "user.out", NULL};
            execv("./spj.out", (char**)command);
        }
        else if (lang == JAVA) {
            const char* command[] = {"./spj.out", "./data.in", "./data.out", "./Main.java", "user.out", NULL};
            execv("./spj.out", (char**)command);
        }
        else {
            exit(-1);
        }
    }
    else if (pid > 0) {
        rusage rusa;
        int status, tt = time_lmt;
        while(1) {
            if (wait4(pid, &status, WNOHANG, &rusa) > 0) {
                break;
            }
            sleep(1);
            --tt;
            if (tt == 0) {
                kill(-pid, SIGKILL);
                usleep(2000);
            }
        }
        //如果运行程序为正常终止
        if (WIFEXITED(status)) {
            int t = WEXITSTATUS(status);
            return t;
        }
        //如果运行程序被信号终止
        else if (WIFSIGNALED(status)) {
            int t = WTERMSIG(status);
            return t;
        }
        return -1;
    }
}
